/**
 * @file       thread_pool.hpp
 * @author     WMQ
 * @date       2024/12/25
 * @brief      线程池
 */

#ifndef THREAD_POOL_HPP
#define THREAD_POOL_HPP

#include "rw_lock.hpp"
#include "thread_func.hpp"
#include <array>
#include <chrono>
#include <condition_variable>
#include <cstdint>
#include <forward_list>
#include <functional>
#include <future>
#include <mutex>
#include <queue>
#include <thread>

#define TVR_BITS (8)   /**< 主动轮位数 */
#define TVR_SIZE (256) /**< 主动轮长度 1 << TVR_BITS */
#define TVN_BITS (6)   /**< 从动轮位数 */
#define TVN_SIZE (64)  /**< 从动轮长度 1 << TVN_BITS */
#define TVR_MASK (255) /**< 主动轮掩码 TVR_SIZE - 1 */
#define TVN_MASK (63)  /**< 从动轮掩码 TVN_SIZE - 1 */

#define TVR_INDEX(v) ((v) & static_cast<uint32_t>(TVR_MASK)) /**< 主动轮当前指针位置 */
#define TVN_INDEX(v, n)                           \
    (((v) >> (TVR_BITS + ((n) - 1) * TVN_BITS)) & \
     static_cast<uint32_t>(TVN_MASK)) /**< 第n个从动轮的当前指针位置 */

#define WHEEL_TOTAL_NUM 5 /**< 从动轮和主动轮总个数 */

/**
 * @brief 线程池
 *
 *   可执行定时任务
 */
class ThreadPool
{
public:
    /**
     * @brief 任务token，一个token控制多个任务，由用户自己创建
     *
     *   通过token控制所有关联任务是否被销毁
     */
    class TaskToken
    {
    public:
        /**
         * @brief 销毁关联的所有任务
         */
        void DestroyTask()
        {
            rw_lock_.WriteLock();
            is_destroy_ = true;
            rw_lock_.WriteUnlock();
        }

    private:
        bool is_destroy_ = false; /**< 是否被销毁 */
        RWLock rw_lock_;          /**< 读写锁 */

        friend ThreadPool;

        void ExecuteTask(const std::function<void()> &task)
        {
            rw_lock_.ReadLock();
            if (!is_destroy_)
            {
                task();
            }
            rw_lock_.ReadUnlock();
        }
    };

    /**
     * @brief 构造函数
     * @param[in] max_size 任务最大数
     * @param[in] thread_num 线程数量
     */
    ThreadPool(uint32_t max_size, uint32_t thread_num) : max_size_(max_size)
    {
        if (0 == thread_num)
        {
            thread_num_ = DEFAULT_THREAD_NUM;
        }
        else
        {
            thread_num_ = thread_num;
        }

        for (uint32_t i = 0; i < thread_num_; ++i)
        {
            threads_.emplace_back([this]() {
                while (!stop_)
                {
                    TaskInfo task_info;
                    if (GetTask(task_info))
                    {
                        task_info.task_token->ExecuteTask(task_info.task_func);
                    }
                }
            });
        }

        start_time_ = std::chrono::steady_clock::now();
        timeout_check_thread_ = std::thread([this]() {
            TimeoutThread();
        });
    }

    /**
     * @brief 析构函数
     */
    ~ThreadPool()
    {
        timeout_check_thread_exit_flag_ = true;
        if (timeout_check_thread_.joinable())
        {
            timeout_check_thread_.join();
        }

        Stop();
        for (uint32_t i = 0; i < thread_num_; ++i)
        {
            if (threads_[i].joinable())
            {
                threads_[i].join();
            }
        }
    }

    /**
     * @brief 向线程池任务队列添加任务
     * @param[in] task_token 任务的token
     * @param[in] task_func  需要执行的任务
     * @return
     *   true  添加成功 \n
     *   false 添加失败 \n
     */
    bool AddTask(const std::shared_ptr<TaskToken> &task_token,
                 const std::function<void()> &task_func)
    {
        TaskInfo task_info;
        task_info.task_token = task_token;
        task_info.task_func = task_func;
        {
            const std::lock_guard<std::mutex> lock(mtx_);
            if (stop_ || task_queue_.size() >= max_size_)
            {
                return false;
            }
            task_queue_.emplace(task_info);
        }
        cond_.notify_one();
        return true;
    }

    /**
     * @brief 向时间轮添加定时任务
     * @param[in] task_token 任务的token
     * @param[in] task_func  需要执行的任务
     * @param[in] timeout    定时时间，单位ms
     * @return
     *   true  添加成功 \n
     *   false 添加失败 \n
     */
    bool AddTask(const std::shared_ptr<TaskToken> &task_token,
                 const std::function<void()> &task_func, uint32_t timeout)
    {
        if (0 == timeout)
        {
            return AddTask(task_token, task_func);
        }

        TimedTaskInfo timed_task_info;
        timed_task_info.task_token = task_token;
        timed_task_info.task_func = task_func;

        const std::lock_guard<std::mutex> lock(time_wheel_mutex_);

        const uint64_t expect_time = timeout + current_index_;
        if (timeout < 256U) // 1 << TVR_BITS
        {
            timed_task_info.current_wheel = 1;
            timed_task_info.position[0] = TVR_INDEX(static_cast<uint32_t>(expect_time));
            main_wheel_0_.at(timed_task_info.position[0]).emplace_front(timed_task_info);
        }
        else if (timeout < 16384U) // 1 << (TVR_BITS + TVN_BITS)
        {
            timed_task_info.current_wheel = 2;
            timed_task_info.position[0] = TVR_INDEX(static_cast<uint32_t>(expect_time));
            timed_task_info.position[1] = TVN_INDEX(static_cast<uint32_t>(expect_time), 1);
            driven_wheel_1_.at(timed_task_info.position[1]).emplace_front(timed_task_info);
        }
        else if (timeout < 1048576U) // 1 << (TVR_BITS + 2 * TVN_BITS)
        {
            timed_task_info.current_wheel = 3;
            timed_task_info.position[0] = TVR_INDEX(static_cast<uint32_t>(expect_time));
            timed_task_info.position[1] = TVN_INDEX(static_cast<uint32_t>(expect_time), 1);
            timed_task_info.position[2] = TVN_INDEX(static_cast<uint32_t>(expect_time), 2);
            driven_wheel_2_.at(timed_task_info.position[2]).emplace_front(timed_task_info);
        }
        else if (timeout < 67108864U) // 1 << (TVR_BITS + 3 * TVN_BITS)
        {
            timed_task_info.current_wheel = 4;
            timed_task_info.position[0] = TVR_INDEX(static_cast<uint32_t>(expect_time));
            timed_task_info.position[1] = TVN_INDEX(static_cast<uint32_t>(expect_time), 1);
            timed_task_info.position[2] = TVN_INDEX(static_cast<uint32_t>(expect_time), 2);
            timed_task_info.position[3] = TVN_INDEX(static_cast<uint32_t>(expect_time), 3);
            driven_wheel_3_.at(timed_task_info.position[3]).emplace_front(timed_task_info);
        }
        else
        {
            timed_task_info.current_wheel = 5;
            timed_task_info.position[0] = TVR_INDEX(static_cast<uint32_t>(expect_time));
            timed_task_info.position[1] = TVN_INDEX(static_cast<uint32_t>(expect_time), 1);
            timed_task_info.position[2] = TVN_INDEX(static_cast<uint32_t>(expect_time), 2);
            timed_task_info.position[3] = TVN_INDEX(static_cast<uint32_t>(expect_time), 3);
            timed_task_info.position[4] = TVN_INDEX(static_cast<uint32_t>(expect_time), 4);
            driven_wheel_4_.at(timed_task_info.position[4]).emplace_front(timed_task_info);
        }

        return true;
    }

    ThreadPool(const ThreadPool &rhs) = delete;
    ThreadPool &operator=(const ThreadPool &rhs) = delete;

private:
    struct TaskInfo
    {
        std::shared_ptr<TaskToken> task_token;
        std::function<void()> task_func;
    };

    struct TimedTaskInfo
    {
        std::shared_ptr<TaskToken> task_token;
        std::function<void()> task_func;                   /**< 执行的函数 */
        std::array<uint32_t, WHEEL_TOTAL_NUM> position {}; /**< 记录此定时器在每个轮中的指针位置 */
        int32_t current_wheel = 0; /**< 记录此定时器当前处在那个轮中 */
    };

    std::vector<std::thread> threads_; /**< 工作线程数组 */
    bool stop_ = false;                /**< 线程池是否停止 */
    std::mutex mtx_;
    std::condition_variable cond_;
    std::queue<TaskInfo> task_queue_;             /**< 任务队列 */
    uint32_t max_size_;                           /**< 任务队列最大任务数 */
    uint32_t thread_num_;                         /**< 工作线程数量 */
    static const uint32_t DEFAULT_THREAD_NUM = 8; /**< 默认工作线程数 */

    std::thread timeout_check_thread_; /**< 用于循环检测定时任务是否超时的线程 */
    bool timeout_check_thread_exit_flag_ = false;      /**< 线程退出标识 */
    std::chrono::steady_clock::time_point start_time_; /**< 起始时间 */
    uint64_t current_index_ = 0;                       /**< 当前时间计数 单位: ms */
    std::mutex time_wheel_mutex_;                      /**< 时间轮资源锁 */
    std::array<std::forward_list<TimedTaskInfo>, TVR_SIZE> main_wheel_0_;   /**< 主动轮0 */
    std::array<std::forward_list<TimedTaskInfo>, TVN_SIZE> driven_wheel_1_; /**< 从动轮1 */
    std::array<std::forward_list<TimedTaskInfo>, TVN_SIZE> driven_wheel_2_; /**< 从动轮2 */
    std::array<std::forward_list<TimedTaskInfo>, TVN_SIZE> driven_wheel_3_; /**< 从动轮3 */
    std::array<std::forward_list<TimedTaskInfo>, TVN_SIZE> driven_wheel_4_; /**< 从动轮4 */

    /**
     * @brief 将某个轮上的某个刻度任务链表上的所有任务移动到下一个轮的某个刻度任务链表
     * @param[in] timed_task_list 刻度任务链表
     * @note
     *   无法移动到下一个轮则超时，将会添加到线程池任务队列
     */
    void MoveTimedTask(std::forward_list<TimedTaskInfo> &timed_task_list)
    {
        while (!timed_task_list.empty())
        {
            bool is_timeout = true;
            TimedTaskInfo &timed_task_info = timed_task_list.front();
            while (0 != (--timed_task_info.current_wheel))
            {
                const uint32_t position =
                    timed_task_info.position.at(timed_task_info.current_wheel - 1);
                if (0 != position)
                {
                    switch (timed_task_info.current_wheel)
                    {
                    case 4:
                        driven_wheel_3_.at(position).emplace_front(timed_task_info);
                        timed_task_list.pop_front();
                        is_timeout = false;
                        break;
                    case 3:
                        driven_wheel_2_.at(position).emplace_front(timed_task_info);
                        timed_task_list.pop_front();
                        is_timeout = false;
                        break;
                    case 2:
                        driven_wheel_1_.at(position).emplace_front(timed_task_info);
                        timed_task_list.pop_front();
                        is_timeout = false;
                        break;
                    case 1:
                        main_wheel_0_.at(position).emplace_front(timed_task_info);
                        timed_task_list.pop_front();
                        is_timeout = false;
                        break;
                    default:
                        break;
                    }
                    break;
                }
            }

            if (is_timeout)
            {
                AddTask(timed_task_info.task_token, timed_task_info.task_func);
                timed_task_list.pop_front();
            }
        }
    }

    /**
     * @brief 检测定时任务是否超时的线程
     * @note
     *   由于steady_clock是一个单调时钟，一但超过其最大值，定时任务将无法超时
     */
    void TimeoutThread()
    {
        while (!timeout_check_thread_exit_flag_)
        {
            ThreadFunc::SleepFor(1);
            const std::lock_guard<std::mutex> lock(time_wheel_mutex_);

            std::chrono::steady_clock::time_point current_time = std::chrono::steady_clock::now();
            while (current_index_ <
                   std::chrono::duration_cast<std::chrono::milliseconds>(current_time - start_time_)
                       .count())
            {
                ++current_index_;

                const uint32_t main_wheel_0_position =
                    TVR_INDEX(static_cast<uint32_t>(current_index_));
                MoveTimedTask(main_wheel_0_.at(main_wheel_0_position));

                if (0 == main_wheel_0_position)
                {
                    const uint32_t driven_wheel_1_position =
                        TVN_INDEX(static_cast<uint32_t>(current_index_), 1);
                    MoveTimedTask(driven_wheel_1_.at(driven_wheel_1_position));

                    if (0 == driven_wheel_1_position)
                    {
                        const uint32_t driven_wheel_2_position =
                            TVN_INDEX(static_cast<uint32_t>(current_index_), 2);
                        MoveTimedTask(driven_wheel_2_.at(driven_wheel_2_position));

                        if (0 == driven_wheel_2_position)
                        {
                            const uint32_t driven_wheel_3_position =
                                TVN_INDEX(static_cast<uint32_t>(current_index_), 3);
                            MoveTimedTask(driven_wheel_3_.at(driven_wheel_3_position));

                            if (0 == driven_wheel_3_position)
                            {
                                const uint32_t driven_wheel_4_position =
                                    TVN_INDEX(static_cast<uint32_t>(current_index_), 4);
                                MoveTimedTask(driven_wheel_4_.at(driven_wheel_4_position));
                            }
                        }
                    }
                }
                current_time = std::chrono::steady_clock::now();
            }
        }
    }

    /**
     * @brief 从线程池任务队列获取一个任务
     * @param[out] task_info 获取的任务
     * @return
     *   true  获取成功 \n
     *   false 获取失败 \n
     */
    bool GetTask(TaskInfo &task_info)
    {
        std::unique_lock<std::mutex> lock(mtx_);
        if (stop_)
        {
            return false;
        }

        while (task_queue_.empty())
        {
            cond_.wait(lock);
            if (stop_)
            {
                return false;
            }
        }

        task_info = std::move(task_queue_.front());
        task_queue_.pop();
        return true;
    }

    /**
     * @brief 停止线程池工作线程
     */
    void Stop()
    {
        {
            const std::lock_guard<std::mutex> lock(mtx_);
            stop_ = true;
        }
        cond_.notify_all();
    }
};
#endif